/*
 * Copyright 2017 - 2025 the original author or authors.
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program. If not, see [https://www.gnu.org/licenses/]
 */

package infra.jdbc.format;

import java.util.HashSet;
import java.util.LinkedList;
import java.util.Locale;
import java.util.StringTokenizer;

import infra.lang.TodayStrategies;

/**
 * Performs formatting of basic SQL statements (DML + query).
 *
 * @author Gavin King
 * @author Steve Ebersole
 * @author <a href="https://github.com/TAKETODAY">Harry Yang</a>
 * @since 4.0 2022/9/12 19:20
 */
public class BasicSQLFormatter implements SQLFormatter {
  private static final boolean SKIP_DDL = TodayStrategies.getFlag("sql.BasicSQLFormatter.skip-ddl", true);

  private static final HashSet<String> BEGIN_CLAUSES = new HashSet<>();
  private static final HashSet<String> END_CLAUSES = new HashSet<>();
  private static final HashSet<String> LOGICAL = new HashSet<>();
  private static final HashSet<String> QUANTIFIERS = new HashSet<>();
  private static final HashSet<String> DML = new HashSet<>();
  private static final HashSet<String> MISC = new HashSet<>();

  static {
    BEGIN_CLAUSES.add("left");
    BEGIN_CLAUSES.add("right");
    BEGIN_CLAUSES.add("inner");
    BEGIN_CLAUSES.add("outer");
    BEGIN_CLAUSES.add("group");
    BEGIN_CLAUSES.add("order");

    END_CLAUSES.add("where");
    END_CLAUSES.add("set");
    END_CLAUSES.add("having");
    END_CLAUSES.add("from");
    END_CLAUSES.add("by");
    END_CLAUSES.add("join");
    END_CLAUSES.add("into");
    END_CLAUSES.add("union");

    LOGICAL.add("and");
    LOGICAL.add("or");
    LOGICAL.add("when");
    LOGICAL.add("else");
    LOGICAL.add("end");

    QUANTIFIERS.add("in");
    QUANTIFIERS.add("all");
    QUANTIFIERS.add("exists");
    QUANTIFIERS.add("some");
    QUANTIFIERS.add("any");

    DML.add("insert");
    DML.add("update");
    DML.add("delete");

    MISC.add("select");
    MISC.add("on");
  }

  private static final String INDENT_STRING = "    ";
  private static final String INITIAL = System.lineSeparator() + INDENT_STRING;

  public static final BasicSQLFormatter INSTANCE = new BasicSQLFormatter();

  private final boolean skipDDL;

  public BasicSQLFormatter() {
    this(SKIP_DDL);
  }

  public BasicSQLFormatter(boolean skipDDL) {
    this.skipDDL = skipDDL;
  }

  @Override
  public String format(String source) {
    if (skipDDL) {
      String lowerCase = source.toLowerCase(Locale.ROOT);
      if (lowerCase.startsWith("create")
              || lowerCase.startsWith("alter table")
              || lowerCase.startsWith("comment on")) {
        return System.lineSeparator() + source;
      }
    }

    return new BasicInternalFormatter().format(source);
  }

  private static final class BasicInternalFormatter {
    boolean beginLine = true;
    boolean afterBeginBeforeEnd;
    boolean afterByOrSetOrFromOrSelect;
    boolean afterOn;
    boolean afterBetween;
    boolean afterInsert;
    int inFunction;
    int parensSinceSelect;
    private final LinkedList<Integer> parenCounts = new LinkedList<>();
    private final LinkedList<Boolean> afterByOrFromOrSelects = new LinkedList<>();

    int indent = 1;

    StringBuilder result = new StringBuilder();
    String lastToken;
    String token;
    String lcToken;

    public String format(String sql) {
      StringTokenizer tokens = new StringTokenizer(
              sql,
              "()+*/-=<>'`\"[]," + WHITESPACE,
              true
      );

      result.append(INITIAL);
      while (tokens.hasMoreTokens()) {
        token = tokens.nextToken();
        lcToken = token.toLowerCase(Locale.ROOT);

        switch (token) {
          case "'" -> {
            String t;
            do {
              t = tokens.nextToken();
              token += t;
            }
            // cannot handle single quotes
            while (!"'".equals(t) && tokens.hasMoreTokens());
          }
          case "\"" -> {
            String t;
            do {
              t = tokens.nextToken();
              token += t;
            }
            while (!"\"".equals(t) && tokens.hasMoreTokens());
          }

          // SQL Server uses "[" and "]" to escape reserved words
          // see SQLServerDialect.openQuote and SQLServerDialect.closeQuote
          case "[" -> {
            String t;
            do {
              t = tokens.nextToken();
              token += t;
            }
            while (!"]".equals(t) && tokens.hasMoreTokens());
          }
        }

        if (afterByOrSetOrFromOrSelect && ",".equals(token)) {
          commaAfterByOrFromOrSelect();
        }
        else if (afterOn && ",".equals(token)) {
          commaAfterOn();
        }

        else if ("(".equals(token)) {
          openParen();
        }
        else if (")".equals(token)) {
          closeParen();
        }

        else if (BEGIN_CLAUSES.contains(lcToken)) {
          beginNewClause();
        }

        else if (END_CLAUSES.contains(lcToken)) {
          endNewClause();
        }

        else if ("select".equals(lcToken)) {
          select();
        }

        else if (DML.contains(lcToken)) {
          updateOrInsertOrDelete();
        }

        else if ("values".equals(lcToken)) {
          values();
        }

        else if ("on".equals(lcToken)) {
          on();
        }

        else if (afterBetween && lcToken.equals("and")) {
          misc();
          afterBetween = false;
        }

        else if (LOGICAL.contains(lcToken)) {
          logical();
        }

        else if (isWhitespace(token)) {
          white();
        }

        else {
          misc();
        }

        if (!isWhitespace(token)) {
          lastToken = lcToken;
        }

      }
      return result.toString();
    }

    private void commaAfterOn() {
      out();
      indent--;
      newline();
      afterOn = false;
      afterByOrSetOrFromOrSelect = true;
    }

    private void commaAfterByOrFromOrSelect() {
      out();
      newline();
    }

    private void logical() {
      if ("end".equals(lcToken)) {
        indent--;
      }
      newline();
      out();
      beginLine = false;
    }

    private void on() {
      indent++;
      afterOn = true;
      newline();
      out();
      beginLine = false;
    }

    private void misc() {
      out();
      if ("between".equals(lcToken)) {
        afterBetween = true;
      }
      if (afterInsert) {
        newline();
        afterInsert = false;
      }
      else {
        beginLine = false;
        if ("case".equals(lcToken)) {
          indent++;
        }
      }
    }

    private void white() {
      if (!beginLine) {
        result.append(" ");
      }
    }

    private void updateOrInsertOrDelete() {
      out();
      indent++;
      beginLine = false;
      if ("update".equals(lcToken)) {
        newline();
      }
      if ("insert".equals(lcToken)) {
        afterInsert = true;
      }
    }

    private void select() {
      out();
      indent++;
      newline();
      parenCounts.addLast(parensSinceSelect);
      afterByOrFromOrSelects.addLast(afterByOrSetOrFromOrSelect);
      parensSinceSelect = 0;
      afterByOrSetOrFromOrSelect = true;
    }

    private void out() {
      result.append(token);
    }

    private void endNewClause() {
      if (!afterBeginBeforeEnd) {
        indent--;
        if (afterOn) {
          indent--;
          afterOn = false;
        }
        newline();
      }
      out();
      if (!"union".equals(lcToken)) {
        indent++;
      }
      newline();
      afterBeginBeforeEnd = false;
      afterByOrSetOrFromOrSelect = "by".equals(lcToken)
              || "set".equals(lcToken)
              || "from".equals(lcToken);
    }

    private void beginNewClause() {
      if (!afterBeginBeforeEnd) {
        if (afterOn) {
          indent--;
          afterOn = false;
        }
        indent--;
        newline();
      }
      out();
      beginLine = false;
      afterBeginBeforeEnd = true;
    }

    private void values() {
      indent--;
      newline();
      out();
      indent++;
      newline();
    }

    private void closeParen() {
      parensSinceSelect--;
      if (parensSinceSelect < 0) {
        indent--;
        parensSinceSelect = parenCounts.removeLast();
        afterByOrSetOrFromOrSelect = afterByOrFromOrSelects.removeLast();
      }
      if (inFunction > 0) {
        inFunction--;
        out();
      }
      else {
        if (!afterByOrSetOrFromOrSelect) {
          indent--;
          newline();
        }
        out();
      }
      beginLine = false;
    }

    private void openParen() {
      if (isFunctionName(lastToken) || inFunction > 0) {
        inFunction++;
      }
      beginLine = false;
      if (inFunction > 0) {
        out();
      }
      else {
        out();
        if (!afterByOrSetOrFromOrSelect) {
          indent++;
          newline();
          beginLine = true;
        }
      }
      parensSinceSelect++;
    }

    private static boolean isFunctionName(String tok) {
      if (tok == null || tok.isEmpty()) {
        return false;
      }

      final char begin = tok.charAt(0);
      final boolean isIdentifier = Character.isJavaIdentifierStart(begin) || '"' == begin;
      return isIdentifier &&
              !LOGICAL.contains(tok) &&
              !END_CLAUSES.contains(tok) &&
              !QUANTIFIERS.contains(tok) &&
              !DML.contains(tok) &&
              !MISC.contains(tok);
    }

    private static boolean isWhitespace(String token) {
      return WHITESPACE.contains(token);
    }

    private void newline() {
      result.append(System.lineSeparator());
      result.append(INDENT_STRING.repeat(Math.max(0, indent)));
      beginLine = true;
    }
  }

}
